import csv
import argparse

# 读取BED文件，返回包含(chromosome, pos, strand)的集合
def read_bed_file(file_path):
    bed_set = set()
    with open(file_path, 'r') as bed_file:
        reader = csv.reader(bed_file, delimiter='\t')
        for row in reader:
            if len(row) >= 3:
                bed_set.add((row[0], int(row[1]), row[2]))
    return bed_set

# 读取TSV文件，并根据正样本和负样本分别进行分析并输出符合条件的记录
def analyze_and_output_tsv(tsv_path, positive_bed, negative_bed, output_path):
    with open(tsv_path, 'r') as tsv_file, open(output_path, 'w', newline='') as output_file:
        reader = csv.reader(tsv_file, delimiter='\t')
        writer = csv.writer(output_file, delimiter='\t')

        # 写入标题行
        #writer.writerow(['chromosome', 'pos', 'strand', 'read_name'])

        next(reader)  # 跳过标题行，如果有的话
        for row in reader:
            if len(row) >= 10:
                chromosome, pos, strand = row[0], int(row[1]), row[2]
                prob_0, prob_1 = float(row[6]), float(row[7])
                read_name = row[4]
                key = (chromosome, pos, strand)

                # 对正样本进行分析
                if key in positive_bed:
                    if prob_1 <= 0.6:
                        writer.writerow([chromosome, pos, strand, read_name])

                # 对负样本进行分析
                if key in negative_bed:
                    if prob_0 <= 0.6:
                        writer.writerow([chromosome, pos, strand, read_name])

def main():
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description="分析TSV文件中与正负样本BED匹配的记录并输出到TSV文件")
    parser.add_argument("positive_bed_path", help="正样本BED文件路径")
    parser.add_argument("negative_bed_path", help="负样本BED文件路径")
    parser.add_argument("tsv_path", help="TSV文件路径")
    parser.add_argument("output_path", help="输出TSV文件路径")
    args = parser.parse_args()

    # 读取正样本和负样本BED文件
    positive_bed = read_bed_file(args.positive_bed_path)
    negative_bed = read_bed_file(args.negative_bed_path)

    # 读取TSV文件并进行分析，结果输出到TSV文件
    analyze_and_output_tsv(args.tsv_path, positive_bed, negative_bed, args.output_path)

if __name__ == "__main__":
    main()
